package com.ververica.cdc.guass.sink.jdbc.derby.table.base;

import com.ververica.cdc.guass.Constants;
import com.ververica.cdc.guass.sink.jdbc.internal.GenericJdbcSinkFunction;
import com.ververica.cdc.guass.sink.jdbc.JdbcTestBase;
import com.ververica.cdc.guass.sink.jdbc.testutils.DatabaseTest;
import com.ververica.cdc.guass.sink.jdbc.testutils.TableManaged;
import com.ververica.cdc.guass.sink.jdbc.testutils.tables.TableRow;
import com.ververica.cdc.guass.source.kafka.data.ChangeEvent;
import com.ververica.cdc.guass.source.kafka.data.NullableValue;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.StateSnapshotContextSynchronousImpl;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkContextUtil;
import org.apache.flink.streaming.util.MockStreamingRuntimeContext;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.catalog.ResolvedSchema;
import org.apache.flink.table.connector.sink.DynamicTableSink;
import org.apache.flink.table.connector.sink.SinkFunctionProvider;
import org.apache.flink.table.data.*;
import org.apache.flink.table.planner.factories.TestValuesTableFactory;
import org.apache.flink.table.runtime.connector.sink.SinkRuntimeProviderContext;
import org.apache.flink.table.types.logical.BigIntType;
import org.apache.flink.table.types.logical.DecimalType;
import org.apache.flink.table.types.logical.IntType;
import org.apache.flink.table.types.logical.LogicalType;
import org.apache.flink.test.util.AbstractTestBase;
import org.apache.flink.types.Row;
import org.apache.flink.types.RowKind;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.charset.StandardCharsets;
import java.sql.Timestamp;
import java.time.format.DateTimeFormatter;
import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

import static com.ververica.cdc.guass.sink.jdbc.testutils.tables.TableBuilder.*;
import static org.apache.flink.table.factories.utils.FactoryMocks.createTableSink;
import static org.assertj.core.api.Assertions.assertThat;

public abstract class JdbcUpdateSinkITCase extends AbstractTestBase implements DatabaseTest {

    private  final Logger log = LoggerFactory.getLogger(this.getClass());

    protected final TableRow upsertOutputTable = createUpsertOutputTable();

    protected final TableRow checkpointOutputTable = createCheckpointOutputTable();

    protected TableRow createCheckpointOutputTable() {
        return tableRow("checkpointTable", field("id", DataTypes.BIGINT().notNull()));
    }


    protected TableRow createUpsertOutputTable() {
        return tableRow(
                "dynamicSinkForUpsert",
                pkField("cnt", DataTypes.BIGINT().notNull()),
                field("lencnt", DataTypes.BIGINT()),
                pkField("cTag", DataTypes.INT().notNull()),
                field("ts", dbType("TIMESTAMP"), DataTypes.TIMESTAMP()));
    }

    protected Map<String, String> getOptions() {
        Map<String, String> options = new HashMap<>();
        options.put("connector", Constants.GAUSS_JDBC_NAME);
        options.put("url", getMetadata().getJdbcUrl());
        options.put("username", getMetadata().getUsername());
        options.put("password", getMetadata().getPassword());
        options.put("table-name", checkpointOutputTable.getTableName());
        options.put("sink.buffer-flush.interval", "0");
        return options;
    }

    @Override
    public List<TableManaged> getManagedTables() {
        return Arrays.asList(
                upsertOutputTable,
                checkpointOutputTable);
    }

    @AfterEach
    void afterEach() {
        TestValuesTableFactory.clearAllData();
    }

    @Test
    protected void testUpdatehangeEventCase() throws Exception {


        // Step 1: Insert initial data via upsert sink
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.getConfig().enableObjectReuse();
        StreamTableEnvironment tEnv = StreamTableEnvironment.create(env);

        String tableName = upsertOutputTable.getTableName();

        tEnv.executeSql(
                upsertOutputTable.getCreateQueryForFlink(
                        getMetadata(),
                        tableName,
                        Arrays.asList(
                                "'sink.buffer-flush.max-rows' = '2'",
                                "'sink.buffer-flush.interval' = '0'",
                                "'sink.max-retries' = '0'")));

        // Insert initial data from testData()
        List<Row> initialData = Arrays.asList(
                Row.of(1L, 10L, 1, Timestamp.valueOf("1970-01-01 00:00:00.001")),
                Row.of(2L, 20L, 2, Timestamp.valueOf("1970-01-01 00:00:00.002")),
                Row.of(3L, 30L, 3, Timestamp.valueOf("1970-01-01 00:00:00.003")));

        // Iterate over the initial data and execute insert statements
        for (Row row : initialData) {
            // Extract fields from the row
            long cnt = row.getFieldAs(0);
            long lencnt = row.getFieldAs(1);
            int cTag = row.getFieldAs(2);
            Timestamp ts = row.getFieldAs(3);

            // Create the INSERT SQL statement with the TIMESTAMP keyword
            String insertSql = String.format("INSERT INTO %s VALUES (%d, %d, %d, TO_TIMESTAMP('%s', 'yyyy-MM-dd HH:mm:ss.SSS'))",
                    tableName, cnt, lencnt, cTag, ts);

            // Execute the SQL statement
            tEnv.executeSql(insertSql).await();

            System.out.println("Executing SQL: " + insertSql);


        }

        List<Row> insertedRows = upsertOutputTable.selectAllTable(getMetadata());

        assertThat(insertedRows.size()).isEqualTo(initialData.size());

        // Step 2: Simulate ChangeEvents with opType "U"
        List<ChangeEvent> changeEvents = new ArrayList<>();

        ChangeEvent update1 = new ChangeEvent();
        update1.setOpType("U");
        update1.setData(new HashMap<String, Object>() {
            {
                put("cnt", 1L);
                put("lencnt", 1000L);
                put("cTag", 1);

            }
        });
        changeEvents.add(update1);
        ChangeEvent update2 = new ChangeEvent();
        update2.setOpType("U");
        update2.setData(new HashMap<String, Object>() {
            {
                put("cnt", 2L);
                put("lencnt", 2000L);
                put("cTag", 2);
                put("ts",  Timestamp.valueOf("1970-01-01 00:00:22.000"));
            }
        });
        changeEvents.add(update2);


        ChangeEvent indert = new ChangeEvent();
        indert.setOpType("I");
        indert.setData(new HashMap<String, Object>() {
            {
                put("cnt", 4L);
                put("lencnt", null);
                put("cTag", 4);
                put("ts",  Timestamp.valueOf("1970-01-01 00:00:44.000"));
            }
        });
        changeEvents.add(indert);

        ChangeEvent delete = new ChangeEvent();
        delete.setOpType("D");
        delete.setData(new HashMap<String, Object>() {
            {
                put("cnt", 3L);
                put("cTag", 3);
            }
        });
        changeEvents.add(delete);

        // Step 2: Convert ChangeEvents to RowData and write to the upsert sink
        ResolvedSchema schema = upsertOutputTable.getTableResolvedSchema2();
        Map<String, String> sinkOptions = getOptions();
        sinkOptions.put("table-name", upsertOutputTable.getTableName());
        DynamicTableSink tableSink = createTableSink(schema, sinkOptions);

        // Create sink function
        SinkRuntimeProviderContext context = new SinkRuntimeProviderContext(false);
        SinkFunctionProvider sinkProvider = (SinkFunctionProvider) tableSink.getSinkRuntimeProvider(context);
        GenericJdbcSinkFunction<RowData> sinkFunction = (GenericJdbcSinkFunction<RowData>) sinkProvider.createSinkFunction();

        // Setup sink function
        sinkFunction.setRuntimeContext(new MockStreamingRuntimeContext(true, 1, 0));
        sinkFunction.setInputType(TypeInformation.of(GenericRowData.class), JdbcTestBase.getExecutionConfig(false));
        sinkFunction.open(new Configuration());

        // Invoke sink function with updates
        for (ChangeEvent event : changeEvents) {
            RowData rowData = convertChangeEventToRowData(event);
            if (rowData != null) {
                sinkFunction.invoke(rowData, SinkContextUtil.forTimestamp(System.currentTimeMillis()));
            }
        }

        // Snapshot state to flush the buffer
        sinkFunction.snapshotState(new StateSnapshotContextSynchronousImpl(1, 1));

        // Close the sink
        sinkFunction.close();

        // Step 4: Verify that the updates have been applied correctly
        List<Row> updatedRows = upsertOutputTable.selectAllTable(getMetadata());

        // Create a map from (cnt, cTag) to Row
        Map<String, Row> rowMap = updatedRows.stream().collect(Collectors.toMap(row -> row.getFieldAs(0) + "-" + row.getFieldAs(2), Function.identity()));

        System.out.println(rowMap);
       // Check that (cnt=1, cTag=1) has updated 'ts'
        Row row1 = rowMap.get("1-1");
        assertThat(row1).isNotNull();
        assertThat(row1.getField(1)). isEqualTo(1000L); // lencnt updated to 1000

        // Check that (cnt=2, cTag=2) has updated 'lencnt' and 'ts'
        Row row2 = rowMap.get("2-2");
        assertThat(row2).isNotNull();
        assertThat(row2.getField(1)).isEqualTo(2000L); // lencnt  updated to 2000
        assertThat(row2.getField(3)).isEqualTo(Timestamp.valueOf("1970-01-01 00:00:22.000")); // ts updated

        // Check that (cnt=3, cTag=3) has deleted
        Row row3 = rowMap.get("3-3");
        assertThat(row3).isNull();


        // Check that (cnt=4, cTag=4) has inderted
        Row row4 = rowMap.get("4-4");
        assertThat(row4).isNotNull();
        assertThat(row4.getField(3)).isEqualTo(Timestamp.valueOf("1970-01-01 00:00:44.000"));// ts updated

        // Print the entire table data
        System.out.println("Updated table data:");
        updatedRows.forEach(System.out::println);

    }





    private Object convertToDataType(Object value, LogicalType logicalType) {
        if (value == null) {
            return null;
        }
        switch (logicalType.getTypeRoot()) {
            case INTEGER:
                return ((Number) value).intValue();
            case BIGINT:
                return ((Number) value).longValue();
            case SMALLINT:
                return ((Number) value).shortValue();
            case TINYINT:
                return ((Number) value).byteValue();
            case DOUBLE:
                return ((Number) value).doubleValue();
            case FLOAT:
                return ((Number) value).floatValue();
            case DECIMAL:
                DecimalType decimalType = (DecimalType) logicalType;
                return DecimalData.fromBigDecimal(new java.math.BigDecimal(value.toString()),
                        decimalType.getPrecision(), decimalType.getScale());
            case VARCHAR:
            case CHAR:
                return StringData.fromString(value.toString());
            case BOOLEAN:
                if (value instanceof Boolean) {
                    return (Boolean) value;
                } else {
                    return Boolean.parseBoolean(value.toString());
                }
            case DATE:
                try {
                    java.time.LocalDate date = java.time.LocalDate.parse(value.toString(), DateTimeFormatter.ISO_DATE);
                    return (int) date.toEpochDay();
                } catch (java.time.format.DateTimeParseException e) {
                    // 尝试其他日期格式
                    try {
                        java.time.LocalDate date = java.time.LocalDate.parse(value.toString(), DateTimeFormatter.ofPattern("yyyy-MM-dd"));
                        return (int) date.toEpochDay();
                    } catch (java.time.format.DateTimeParseException ex) {
                        return null;
                    }
                }
            case TIME_WITHOUT_TIME_ZONE:
                try {
                    java.time.LocalTime time = java.time.LocalTime.parse(value.toString(), DateTimeFormatter.ISO_TIME);
                    return (int) (time.toNanoOfDay() / 1_000_000); // 转换为毫秒
                } catch (java.time.format.DateTimeParseException e) {
                    // 尝试其他时间格式
                    try {
                        java.time.LocalTime time = java.time.LocalTime.parse(value.toString(), DateTimeFormatter.ofPattern("HH:mm:ss"));
                        return (int) (time.toNanoOfDay() / 1_000_000);
                    } catch (java.time.format.DateTimeParseException ex) {
                        return null;
                    }
                }
            case TIMESTAMP_WITHOUT_TIME_ZONE:
                try {
                    DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss[.S]");
                    java.time.LocalDateTime dateTime = java.time.LocalDateTime.parse(value.toString(), formatter);
                    return TimestampData.fromLocalDateTime(dateTime);
                } catch (java.time.format.DateTimeParseException e) {
                    // 尝试其他时间戳格式
                    try {
                        java.time.LocalDateTime dateTime = java.time.LocalDateTime.parse(value.toString(), DateTimeFormatter.ISO_DATE_TIME);
                        return TimestampData.fromLocalDateTime(dateTime);
                    } catch (java.time.format.DateTimeParseException ex) {
                        return null;
                    }
                }
            case BINARY:
            case VARBINARY:
                if (value instanceof byte[]) {
                    return value;
                } else if (value instanceof String) {
                    return ((String) value).getBytes(StandardCharsets.UTF_8);
                } else {
                    return null;
                }
            default:
                throw new UnsupportedOperationException("Unsupported data type: " + logicalType);
        }
    }

    /**
     * Helper method to convert ChangeEvent to RowData.
     */
    private RowData convertChangeEventToRowData(ChangeEvent event) {
        String opType = event.getOpType();
        RowKind rowKind;
        HashMap<String, Object> dataMap = new HashMap<>();
        dataMap.putAll(event.getData());

        switch (opType) {
            case "I":
                rowKind = RowKind.INSERT;
                break;
            case "U":
                rowKind = RowKind.UPDATE_AFTER;
                break;
            case "D":
                rowKind = RowKind.DELETE;
                break;
            default:
                // Unknown operation type, skip or throw exception
                return null;
        }

        List<String> fieldNames = Arrays.asList("cnt", "lencnt", "cTag", "ts");

        List<LogicalType> fieldTypes = Arrays.asList(
                new BigIntType(),
                new BigIntType(),
                new IntType(),
                new org.apache.flink.table.types.logical.TimestampType(3)
        );

        // 创建具有指定 RowKind 的 CustomRowData
        GenericRowData rowData = new GenericRowData(rowKind, fieldNames.size());


        for (int i=0;i<fieldNames.size();i++) {
            // 检查字段值是否不为空
            String  fieldName = fieldNames.get(i);
            LogicalType fieldType = fieldTypes.get(i);
            if (dataMap.containsKey(fieldName)) {
                Object fieldValue = dataMap.get(fieldName);
                Object convertedValue = convertToDataType(fieldValue, fieldType);
                rowData.setField(i, new NullableValue(convertedValue, true));
            } else {
                // 使用一个特殊的 NullableValue 对象表示字段不存在
                rowData.setField(i, new NullableValue(null, false));
            }

        }

        return rowData;
    }

}
